from Trader import Trader
import numpy as np
from dayCalculation import *
import math

class CMTrader(Trader):
    def __init__(self, initial_money=0,  tax_ratio=0.001, max_stock_num = 10,stocks_info = {}):
        Trader.__init__(self, initial_money,  tax_ratio, max_stock_num , stocks_info)
        self.stockpool_price = {}

    def stock_price(self,date,internal_days):
        trade_days = list(self.stocks_info[list(self.stocks_info.keys())[0]].keys())
        date_index = trade_days.index(date)
        pre_day = [trade_days[date_index - i] for i in range(internal_days + 1)]
        stock_price = {}
        for stock in self.stocks_info:
            prices = [self.stocks_info[stock][pre_day[i]]['close'] for i in range(0,len(pre_day))]
            stock_price[stock] = prices
        return stock_price

    def stock_earn(self,date,internal_days):
        trade_days = list(self.stocks_info[list(self.stocks_info.keys())[0]].keys())
        date_index = trade_days.index(date)
        pre_day = [trade_days[date_index-i] for i in range(internal_days+1)]
        stock_earn = {}
        for stock in self.stocks_info:
            prices = [self.stocks_info[stock][pre_day[i]]['close'] for i in range(0,len(pre_day))]
            stock_earn[stock] = []
            for k in range(len(prices)-1):
                price1 = prices[k]
                price2 = prices[k+1]
                if not math.isnan(price1) and not math.isnan(price2):
                    stock_earn[stock].append(math.log(price1/price2))
                else:
                    stock_earn[stock].append(0)
        return stock_earn
        
    def strategy(self,date,internal_days):
        down_stocks = []
        up_stocks = []
        stockpool_grossearn = {}
        buy_num = self.max_stock_num
        stock_earn = self.stock_earn(date,internal_days)
        stock_price = self.stock_price(date,internal_days)
        for stock in self.stocks_info:
            earn_rate = stock_earn[stock]

            gross_earn = sum(earn_rate)
            if stock in self.stockpool_price:
                stockpool_grossearn[stock] = gross_earn
            down_day = sum([i<0 for i in earn_rate])
            up_day = sum([i > 0 for i in earn_rate])
            if down_day == len(earn_rate) or gross_earn<-0.05:
                down_stocks.append(stock)

            if up_day == len(earn_rate) or gross_earn>0.05:
                up_stocks.append(stock)
        for stock in down_stocks:
            if stock in self.stockpool_price:
                self.sell(date, stock_name=stock, share=self.stock_shares[stock])
        for stock in self.stockpool_price:
            initial_price = self.stockpool_price[stock]
            current_price = stock_price[stock][0]
            if math.log(current_price/initial_price)>0.2:
                self.sell(date, stock_name=stock, share=self.stock_shares[stock])
        for stock in self.stock_shares:
            if self.stock_shares[stock] > 0:
                buy_num -= 1
        while buy_num < len(up_stocks):
            max_key = max(stockpool_grossearn, key =stockpool_grossearn.get)
            self.sell(date,stock_name=max_key,share = self.stock_shares[max_key])
            buy_num += 1
        for stock in up_stocks:
            money = self.cash / len(up_stocks)
            self.buy(date, stock_name=stock, money=money)
        self.net_worth = self.get_worth(date)







